import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from python_ai.common.xcommon import sep

pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False


def x_to_key(x, x_v, y):
    return x, x_v, y


def x_get_key(P, k):
    x, x_v, y = k
    if k in P.keys():
        return P[k]
    else:
        return P[(x, None, y)]


class NaiveBayesImplWithSmooth(object):

    def __init__(self, lam=1):
        self.lam_ = lam

    def fit(self, df):
        P = {}

        m = len(df)
        y_counts_series = df.iloc[:, -1].value_counts()
        y_values = y_counts_series.index  # all target values
        y_values_len = len(y_values)
        y_p_series = (y_counts_series + self.lam_) / (m + y_values_len * self.lam_)  # P of each target value

        x_labels = df.columns[:-1]  # all features' name

        for y_v in y_values:
            df2 = df[df.iloc[:, -1] == y_v]
            len_of_this_y_v = y_counts_series[y_v]

            for x in x_labels:
                x_count_series = df2[x].value_counts()
                x_values = x_count_series.index
                x_values_len = len(x_values)

                x_p_series = (x_count_series + self.lam_) / (len_of_this_y_v + x_values_len * self.lam_)

                k = x_to_key(x, None, y_v)
                P[k] = self.lam_ / (len_of_this_y_v + x_values_len * self.lam_)
                for x_v in x_values:
                    k = x_to_key(x, x_v, y_v)
                    P[k] = x_p_series[x_v]

        print(y_p_series)  # tmp
        print(P)  # tmp

        self.P_ = P
        self.y_p_series_ = y_p_series
        self.x_labels_ = x_labels
        self.y_values_ = y_values

    def predict(self, X):
        result = []
        for y_v in self.y_values_:
            r = self.y_p_series_[y_v]
            for i, x_v in enumerate(X):
                k = x_to_key(self.x_labels_[i], x_v, y_v)
                v = x_get_key(self.P_, k)
                r *= v
            result.append([y_v, r])
        result = np.array(result)
        idx_sorted = result[:, 1].argsort()[-1::-1]
        result = result[idx_sorted]
        return result[0, 0], result


if '__main__' == __name__:
    sep('load')
    df = pd.read_csv('bayes_lihang.txt', header=0)
    m = len(df)
    print(m)
    print(df[:5])

    model = NaiveBayesImplWithSmooth()
    model.fit(df)

    X = [2, 'S']
    result = model.predict(X)
    print(result)

    X = [2, 'N']
    result = model.predict(X)
    print(result)

    X = [200, 'N']
    result = model.predict(X)
    print(result)
